///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file configuration.cpp
///
/// \brief   implementation for the Configuration class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#include <vector>
#include <cassert>
#include <numeric>
#include <iostream>
#include <cstdio>

#include <src/configuration.h>
#include <src/indexing.h>

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Constructs the configuration object
///
/// \param[in]      ns       the number of sites
/// \param[in]      to       the total lattice occupation
/// \param[in]      cps      the number of configurations on a site
///
///////////////////////////////////////////////////////////////////////////////////////////////////
eds::Configuration::Configuration(const int ns, const int to, const int cps)
  : _ns(ns),
    _to(to),
    _cps(cps),
    _occ_vec(ns, 0),
    _occurances(cps, 0)
{

//  // get the number of configurations, using the occurances vector as workspace
//  _nc = eds::n_integer_solutions(ns, to, cps-1, &_occurances.at(0));
//
//  // reset the occurances vector to zero
//  std::fill(_occurances.begin(), _occurances.end(), 0);

  // reset to the first valid configuration
  this->reset();

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Resets the configuration
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void eds::Configuration::reset() {

  // reset the finished flag
  _finished = false;

  // starting with the largest occupation, set each occupation's occurance to its smallest
  // possible value
  {
    int rem_occ = _to;
    int rem_sites = _ns;
    for (int i = _cps - 1; i > 0; i--) {
      _occurances[i] = 0;
      while( rem_occ - i * _occurances[i] - (i-1) * (rem_sites - _occurances[i]) > 0 )
        _occurances[i]++;
      rem_sites -= _occurances[i];
      rem_occ -= i * _occurances[i];
    }
    assert( rem_occ == 0 );
    assert( rem_sites >= 0 );
  }

  // sort the occupations into the occupation number vector in descending order
  {
    int pos = 0;
    for (int i = _cps - 1; i > 0; i--)
      for (int j = 0; j < _occurances[i]; j++)
        _occ_vec[pos++] = i;
    for ( ; pos < _occ_vec.size(); pos++)
      _occ_vec[pos] = 0;
  }

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Increments to the next valid configuration
///
///////////////////////////////////////////////////////////////////////////////////////////////////
void eds::Configuration::operator++() {

  // loop over the different particle ranks
  bool finished_shifting = true;
  int r = 1;
  while ( r < _cps ) {

    // shift right
    finished_shifting = false;
    bool flag = false;
    int accumulation = 0;
    for (int i = _occ_vec.size() - 1; i >= 0; i--) {
      
      if (_occ_vec[i] == r)
        accumulation++;

      if ( accumulation == _occurances[r] && !flag ) {
        finished_shifting = true;
        break;
      }

      if (_occ_vec[i] < r)
        flag = true;

      if (_occ_vec[i] == r && flag) {
        for (int k = 0; k < i; k++)
          if (_occ_vec[k] < r)
            _occ_vec[k] = -1;
        _occ_vec[i] = -1;
        //accumulation--;
        int pos = i+1;
        for ( ; accumulation > 0; pos++) {
          assert( pos < _occ_vec.size() );
          if (_occ_vec[pos] <= r) {
            _occ_vec[pos] = r;
            accumulation--;
          }
        }
        assert( pos <= _occ_vec.size() );
        for ( ; pos < _occ_vec.size(); pos++)
          if (_occ_vec[pos] <= r)
            _occ_vec[pos] = -1;
        break;
      }

    }

    // if we are not finished shifting at this particle rank, then the increment is complete
    if (!finished_shifting) break;

    // otherwise, we re-distribute the next particle rank in the next iteration of this loop
    r++;

  }

  // if we were are finished shifting, move to the next distribution of sites among particle ranks
  if (finished_shifting) {

    // set the occupation vector to -1 flags
    for (int i = 0; i < _occ_vec.size(); i++)
      _occ_vec[i] = -1;

    // get the remaining occupation
    int rem_occ = _to;
    int rem_sites = _ns;
    for (int s = _cps-1; s > 1; s--) {
      rem_sites -= _occurances[s];
      rem_occ -= s * _occurances[s];
    }

    // redistribute
    int s = 2;
    for ( ; s < _cps; s++) {
      // try adding another occupancy of this rank
      if ( rem_occ - s >= 0 && rem_sites - 1 >= 0 ) {
        _occurances[s]++;
        rem_sites--;
        rem_occ -= s;
        break;
      }
      // if that doesn't work, go up to the next particle rank
      rem_sites += _occurances[s];
      rem_occ += s * _occurances[s];
    }
    //std::cout << s << std::endl;
    if (s == _cps) {
      _finished = true;
      return;
    }
    for (s = s - 1; s > 0; s--) {
      _occurances[s] = 0;
      while( rem_occ - s * _occurances[s] - (s-1) * (rem_sites - _occurances[s]) > 0 )
        _occurances[s]++;
      rem_sites -= _occurances[s];
      rem_occ -= s * _occurances[s];
    }
    assert( rem_occ == 0 );
    assert( rem_sites >= 0 );

    // make sure that r is maximum so that all particle rank's arrangements will be reset
    assert( r == _cps );

  }

  //printf("  %10s    ", " ");
  //for (int i = 0; i < _occ_vec.size(); i++)
  //  printf(" %2i", _occ_vec.at(i));
  //printf("    %2i    %2i", r, _occurances[1]);
  //std::cout << std::endl;

  // reset the lower particle ranks now that we have shifted or moved to a new distribution
  {
    int accumulation = 0;
    for (int pos = 0; pos < _occ_vec.size(); pos++) {
      if (_occ_vec[pos] == -1) {
        while ( accumulation == 0 && r > 0 ) {
          r--;
          accumulation = _occurances[r];
        }
        _occ_vec[pos] = r;
        accumulation--;
      }
    }
    //assert( r == 0 );
  }

  //printf("  %10s    ", " ");
  //for (int i = 0; i < _occ_vec.size(); i++)
  //  printf(" %2i", _occ_vec.at(i));
  //std::cout << std::endl;

  // check that the total occupation remains correct
  assert( std::accumulate(_occ_vec.begin(), _occ_vec.end(), 0) == _to );

  //assert( 1 == 0 );

}

///////////////////////////////////////////////////////////////////////////////////////////////////
/// \brief Assignment operator
///
///////////////////////////////////////////////////////////////////////////////////////////////////
eds::Configuration & eds::Configuration::operator=(const eds::Configuration & other) {

  if (this == &other) return *this;

  if (_ns != other._ns)
    _occ_vec.resize(other._ns);

  if (_cps != other._cps)
    _occurances.resize(other._cps);

  _ns = other._ns;
  _to = other._to;
  _cps = other._cps;
  _finished = other._finished;
  for (int i = 0; i < _ns; i++)
    _occ_vec[i] = other._occ_vec[i];
  for (int i = 0; i < _cps; i++)
    _occurances[i] = other._occurances[i];

  return *this;

}

